import random
import time
import numpy

file = '/Users/apple/Downloads/solverResultData/1001_1000_fdat/1001_1000.fdat'
file2 = '/Users/apple/Downloads/solverResultData/101_1000_fdat/101_1000.fdat'

shape = (1001, 1, 1000, 1000)

def read_with_no_seek():
    start = time.time()
    with open(file, 'rb') as f:
        f.seek(18 + 2001*4)
        f.read(10000 * 4)
    print('read_with_no_seek use time: ', time.time() - start)

def read_and_seek():
    start = time.time()
    buffer = b''
    with open(file, 'rb') as f:
        f.seek(18 + 2001*4)
        for i in range(10000):
            buffer += f.read(4)
            f.seek(4, 1)

    print('read_and_seek use time: ', time.time() - start)

def read_by_mmap():
    start = time.time()

    a = numpy.memmap(
        filename=file, 
        dtype=numpy.float32,
        mode='r', 
        offset=18 + 2001*4, 
        shape=(12, 1001, 100, 100, 100), 
        order='C'
    )
    import random
    index = random.randint(0, 99)
    index1 = random.randint(0, 11)
    slices = {
        'y  * z': lambda a: a[index1, index, index, :, :],
        'x  * z': lambda a: a[index1, index, :, index, :],
        'x  * y': lambda a: a[index1, index, :, :, index],
        'freq*z': lambda a: a[index1, :, index, index, :],
        'freq*y': lambda a: a[index1, :, index, :, index],
        'freq*x': lambda a: a[index1, :, :, index, index],
    }
    for k,v in slices.items():
        b = numpy.square(v(a)) - 0.1
        end = time.time()
        print(k, ' use time: ', end - start)
    
    del a

def read_by_mmap_1():
    start = time.time()

    a = numpy.memmap(
        filename=file, 
        dtype=numpy.float32,
        mode='r', 
        offset=18 + 2001*4, 
        shape=(12, 1001, 1, 1000, 1000), 
        order='C'
    )
    import random
    index = random.randint(0, 99)
    index1 = random.randint(0, 11)
    slices = {
        'y  * z': lambda a: a[index1, index, 0, :, :],
        'x  * z': lambda a: a[index1, index, :, index, :],
        'x  * y': lambda a: a[index1, index, :, :, index],
        'freq*z': lambda a: a[index1, :, 0, index, :],
        'freq*y': lambda a: a[index1, :, 0, :, index],
        'freq * x': lambda a: a[2, :, :, 91, 91],
    }
    for k,v in slices.items():
        b = numpy.square(v(a)) - 0.1
        end = time.time()
        print(k, ' use time: ', end - start)
        start = end
    del a

def read_by_mmap_2():
    start = time.time()

    a = numpy.memmap(
        filename=file2, 
        dtype=numpy.float32,
        mode='r', 
        offset=18 + 2102*4, 
        shape=(12, 101, 1, 1000, 1000), 
        order='C'
    )
    import random
    index = random.randint(0, 99)
    index1 = random.randint(0, 11)
    slices = {
        'y  * z': lambda a: a[index1, index, 0, :, :],
        'x  * z': lambda a: a[index1, index, :, index, :],
        'x  * y': lambda a: a[index1, index, :, :, index],
        'freq*z': lambda a: a[index1, :, 0, index, :],
        'freq*y': lambda a: a[index1, :, 0, :, index],
        'freq*x': lambda a: a[index1, :, :, index, index],
    }
    for k,v in slices.items():
        b = numpy.square(v(a)) - 0.1
        end = time.time()
        print(k, ' use time: ', end - start)
        start = end
    del a

def read_by_mmap_3():
    start = time.time()

    a = numpy.memmap(
        filename=file, 
        dtype=numpy.float32,
        mode='r', 
        offset=18 + 100201*4, 
        shape=(12, 100, 1, 1000, 1000), 
        order='C'
    )
    import random
    index = random.randint(0, 99)
    index1 = random.randint(0, 11)
    slices = {
        'y  * z': lambda a: a[index1, index, 0, :, :],
        'x  * z': lambda a: a[index1, index, :, index, :],
        'x  * y': lambda a: a[index1, index, :, :, index],
        'freq*z': lambda a: a[index1, :, 0, index, :],
        'freq*y': lambda a: a[index1, :, 0, :, index],
        'freq*x': lambda a: a[index1, :, :, index, index],
    }
    for k,v in slices.items():
        b = numpy.square(v(a)) - 0.1
        end = time.time()
        print(k, ' use time: ', end - start)
        start = end
    del a

def create_h5py():
    import h5py
    with h5py.File(file+'.h5', 'w') as f:
        for i in range(3):
            a = numpy.memmap(
                filename=file, 
                dtype=numpy.float32,
                mode='r', 
                offset=18 + 3002*4 + i*1001*1000*1000*4, 
                shape=(1001, 1, 1000, 1000), 
                order='C'
            )
            f.create_dataset(str(i), 
                shape=(1001, 1, 1000, 1000), 
                dtype=numpy.float32,
                data=a
            )

def read_by_h5py():
    start = time.time()
    import h5py
    with h5py.File(file+'.h5', mode='r') as f:
        a = f['1'] 
        index = random.randint(0, 999)
        slices = {
            'y  * z': lambda a: a[index, 0, :, :],
            'x  * z': lambda a: a[index, :, index, :],
            'x  * y': lambda a: a[index, :, :, index],
            'freq*z': lambda a: a[:, 0, index, :],
            'freq*y': lambda a: a[:, 0, :, 199],
            'freq*x': lambda a: a[:, :, index, index],
        }
        for k,v in slices.items():
            b = numpy.square(v(a)) - 0.1
            end = time.time()
            print(b.shape)
            print(k, ' use time: ', end - start)
            start = end         
        # del a

if __name__ == '__main__':
    # 1. 验证seek 的效率， seek 越多越慢
    # 2. 研究如何减少 seek次数，  大文件存两份，其中一份颠倒排列顺序
    # 3. h5py  freq * y 比numpy 快三倍， 但还是不够快

    # read_with_no_seek()
    # read_and_seek()
    # read_by_mmap()
    # read_by_mmap_1()
    # read_by_mmap_2()
    read_by_mmap_3()
    # create_h5py()
    # read_by_h5py()